#include "kernel_operator.h"
#include "float.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

class KernelThreeNN
{
public:
    __aicore__ inline KernelThreeNN() {}
    __aicore__ inline void Init(GM_ADDR xyz1, GM_ADDR xyz2, 
                                GM_ADDR dist, GM_ADDR indices,
                                uint32_t ALIGN_NUM, uint32_t ub_data_num,
                                uint32_t b, uint32_t m, uint32_t n, 
                                uint32_t batchDataNum, uint32_t batchDataNumAligned, 
                                uint32_t core_size, uint32_t core_remain)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        uint32_t blockIdx = GetBlockIdx();

        this->coreDataNum = core_size + (GetBlockNum() == GetBlockIdx() + 1 ? core_remain : 0);
        this->tileDataNum = ub_data_num;
        this->startPointer = batchDataNum * GetBlockIdx();
        this->b = b;
        this->m = m;
        this->n = n;
        this->batchDataNum = batchDataNumAligned;

        xyz1Gm.SetGlobalBuffer((__gm__ DTYPE_XYZ1 *)xyz1, b * n * 3);
        xyz2Gm.SetGlobalBuffer((__gm__ DTYPE_XYZ1 *)xyz2 + this->startPointer, this->coreDataNum);
        dGm.SetGlobalBuffer((__gm__ DTYPE_DIST *)dist, b * n * 3);
        iGm.SetGlobalBuffer((__gm__ DTYPE_INDICES *)indices, b * n * 3);



        this->tileNum = this->coreDataNum / this->tileDataNum + (this->coreDataNum % this->tileDataNum > 0);
        // printf("coreDataNum: %d\ttileNum: %d\n", coreDataNum, tileNum);

        pipe.InitBuffer(xyz0Queue, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_XYZ1));
        pipe.InitBuffer(xyz1Queue, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_XYZ1));
        pipe.InitBuffer(x2Queue, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_XYZ1));
        pipe.InitBuffer(y2Queue, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_XYZ1));
        pipe.InitBuffer(z2Queue, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_XYZ1));
        pipe.InitBuffer(tmpBuffer, this->tileDataNum * sizeof(DTYPE_XYZ1));

        // pipe.InitBuffer(dQueue, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_DIST));
        // pipe.InitBuffer(iQueue, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_INDICES));
        // pipe.InitBuffer(bestQueue, BUFFER_NUM, 32);
        // pipe.InitBuffer(bestiQueue, BUFFER_NUM, 32);
        pipe.InitBuffer(bestBuf, 32);
        pipe.InitBuffer(bestiBuf, 32);

    }

    __aicore__ inline void Process()
    {
        for (int32_t i = 0; i < b; i++)
        {
            for (int32_t j = 0; j < n; j++)
            {
                int32_t xyz0Idx = i * n * 3 + j * 3;
                int32_t loopCount = tileNum;

                InitBestBuf();
                for (int32_t k = 0; k < loopCount; k++)
                {
                    int32_t length = (k < loopCount - 1) ? tileDataNum : coreDataNum - (loopCount - 1) * tileDataNum;
                    // printf("loopCount: %d\n", loopCount);
                    CopyIn(i, k, length);
                    Compute(xyz0Idx, k, length / 3);
                    // Compute(6, k, length / 3);
                    // CopyOut(i, k, length);
                }
            }
        }
    }

    __aicore__ inline void test()
    {
        LocalTensor<DTYPE_XYZ1> x2Local = x2Queue.AllocTensor<DTYPE_XYZ1>();
        LocalTensor<DTYPE_XYZ1> y2Local = y2Queue.AllocTensor<DTYPE_XYZ1>();
        LocalTensor<DTYPE_XYZ1> z2Local = z2Queue.AllocTensor<DTYPE_XYZ1>();

        LocalTensor<DTYPE_DIST> dLocal = dQueue.AllocTensor<DTYPE_XYZ1>();
        LocalTensor<DTYPE_INDICES> iLocal = iQueue.AllocTensor<DTYPE_INDICES>();

        
        DataCopy(x2Local, xyz2Gm[0], 96);

        for (int32_t ii = 0; ii < 96 / 3; ii++)
        {
            x2Local.SetValue(ii, xyz2Gm.GetValue(ii * 3 + 0));
            y2Local.SetValue(ii, xyz2Gm.GetValue(ii * 3 + 1));
            z2Local.SetValue(ii, xyz2Gm.GetValue(ii * 3 + 2));
        }
        
        LocalTensor<DTYPE_XYZ1> tLocal = tmpBuffer.Get<DTYPE_XYZ1>();

        uint32_t length = 32;
        for (int i = 0; i < 32; i++)
        {
            tLocal.SetValue(i, xyz1Gm.GetValue(0));
        }
        Sub(x2Local, x2Local, tLocal, length);
        Mul(x2Local, x2Local, x2Local, length);
        
        for (int i = 0; i < 32; i++)
        {
            tLocal.SetValue(i, xyz1Gm.GetValue(1));
        }
        Sub(y2Local, y2Local, tLocal, length);
        Mul(y2Local, y2Local, y2Local, length);
        Add(x2Local, x2Local, y2Local, length);

        for (int i = 0; i < 32; i++)
        {
            tLocal.SetValue(i, xyz1Gm.GetValue(2));
        }
        DataCopy(dGm, tLocal, 16);
        Sub(z2Local, z2Local, tLocal, length);
        Mul(z2Local, z2Local, z2Local, length);
        Add(x2Local, x2Local, z2Local, length);


        LocalTensor<DTYPE_XYZ1> bLocal = bestQueue.AllocTensor<DTYPE_XYZ1>();
        LocalTensor<DTYPE_INDICES> biLocal = bestiQueue.AllocTensor<DTYPE_INDICES>();

        // bLocal.SetValue(0, (DTYPE_XYZ1)1e38);
        // bLocal.SetValue(1, (DTYPE_XYZ1)1e38);
        // bLocal.SetValue(2, (DTYPE_XYZ1)1e38);
        // biLocal.SetValue(0, (DTYPE_INDICES)0);
        // biLocal.SetValue(1, (DTYPE_INDICES)0);
        // biLocal.SetValue(2, (DTYPE_INDICES)0);

        // for (int32_t i = 0; i < 32; i++)
        // {
        //     if (x2Local.GetValue(i) < bLocal.GetValue(0))
        //     {
        //         bLocal.SetValue(1, bLocal.GetValue(0));
        //         bLocal.SetValue(2, bLocal.GetValue(1));
        //         biLocal.SetValue(1, biLocal.GetValue(0));
        //         biLocal.SetValue(2, biLocal.GetValue(1));
        //         bLocal.SetValue(0, x2Local.GetValue(i));
        //         biLocal.SetValue(0, (int32_t)i);
        //     }
        //     else if (x2Local.GetValue(i) < bLocal.GetValue(1))
        //     {
        //         bLocal.SetValue(2, bLocal.GetValue(1));
        //         biLocal.SetValue(2, biLocal.GetValue(1));
        //         bLocal.SetValue(1, x2Local.GetValue(i));
        //         biLocal.SetValue(1, (int32_t)i);
        //     }
        //     else if (x2Local.GetValue(i) < bLocal.GetValue(2))
        //     {
        //         bLocal.SetValue(2, x2Local.GetValue(i));
        //         biLocal.SetValue(2, (int32_t)i);
        //     }
            
        //     dLocal.SetValue(0, bLocal.GetValue(0));
        //     dLocal.SetValue(1, bLocal.GetValue(1));
        //     dLocal.SetValue(2, bLocal.GetValue(2));
        //     iLocal.SetValue(0, biLocal.GetValue(0));
        //     iLocal.SetValue(1, biLocal.GetValue(1));
        //     iLocal.SetValue(2, biLocal.GetValue(2));
            
        // }

        // DataCopy(dGm, dLocal, 8);
        // DataCopy(iGm, iLocal, 8);


        x2Queue.FreeTensor(x2Local);
        y2Queue.FreeTensor(y2Local);
        z2Queue.FreeTensor(z2Local);

        dQueue.FreeTensor(dLocal);
    }

private:
    __aicore__ inline void CopyIn(int32_t batchIdx, int32_t progress, uint32_t length)
    {
        LocalTensor<DTYPE_XYZ1> x2Local = x2Queue.AllocTensor<DTYPE_XYZ1>();
        LocalTensor<DTYPE_XYZ1> y2Local = y2Queue.AllocTensor<DTYPE_XYZ1>();
        LocalTensor<DTYPE_XYZ1> z2Local = z2Queue.AllocTensor<DTYPE_XYZ1>();

        for (int32_t ii = 0; ii < length / 3; ii++)
        {
            x2Local.SetValue(ii, xyz2Gm.GetValue(batchIdx * m * 3 + progress * tileDataNum + ii * 3 + 0));
            y2Local.SetValue(ii, xyz2Gm.GetValue(batchIdx * m * 3 + progress * tileDataNum + ii * 3 + 1));
            z2Local.SetValue(ii, xyz2Gm.GetValue(batchIdx * m * 3 + progress * tileDataNum + ii * 3 + 2));
        }

        // printf("length: %d\n", length);

        // DataCopy(dGm[progress * tileDataNum / 3], x2Local, 24);

        x2Queue.EnQue(x2Local);        
        y2Queue.EnQue(y2Local);
        z2Queue.EnQue(z2Local);  

    }

    __aicore__ inline void Compute(int32_t idx, int32_t progress, uint32_t length)
    {
        LocalTensor<DTYPE_XYZ1> x2Local = x2Queue.DeQue<DTYPE_XYZ1>();
        LocalTensor<DTYPE_XYZ1> y2Local = y2Queue.DeQue<DTYPE_XYZ1>();
        LocalTensor<DTYPE_XYZ1> z2Local = z2Queue.DeQue<DTYPE_XYZ1>();

        // LocalTensor<DTYPE_DIST> bLocal = bestQueue.DeQue<DTYPE_DIST>();
        // LocalTensor<DTYPE_INDICES> biLocal = bestiQueue.DeQue<DTYPE_INDICES>();
        LocalTensor<DTYPE_XYZ1> bLocal = bestBuf.Get<DTYPE_XYZ1>();
        LocalTensor<DTYPE_INDICES> biLocal = bestiBuf.Get<DTYPE_INDICES>();

        // LocalTensor<DTYPE_DIST> dLocal = dQueue.AllocTensor<DTYPE_DIST>();
        // LocalTensor<DTYPE_INDICES> iLocal = iQueue.AllocTensor<DTYPE_INDICES>();
    
        LocalTensor<DTYPE_XYZ1> tLocal = tmpBuffer.Get<DTYPE_XYZ1>();
        // Duplicate(tLocal, xyz1Gm.GetValue(idx + 0), length);
        Muls(tLocal, tLocal, (float)0.0, length);
        Adds(tLocal, tLocal, xyz1Gm.GetValue(idx + 0), length);
        // for (int i = 0; i < length; i++)
        // {
        //     tLocal.SetValue(i, xyz1Gm.GetValue(idx + 0));
        // }
        // DataCopy(dGm[32], tLocal, 16);
        Sub(x2Local, x2Local, tLocal, length);
        Mul(x2Local, x2Local, x2Local, length);
        // DataCopy(dGm[0], x2Local, 16);
        
        // Duplicate(tLocal, xyz1Gm.GetValue(idx + 1), length);
        Muls(tLocal, tLocal, (float)0.0, length);
        Adds(tLocal, tLocal, xyz1Gm.GetValue(idx + 1), length);
        // for (int i = 0; i < length; i++)
        // {
        //     tLocal.SetValue(i, xyz1Gm.GetValue(idx + 1));
        // }
        // DataCopy(dGm[16], tLocal, 16);
        Sub(y2Local, y2Local, tLocal, length);
        Mul(y2Local, y2Local, y2Local, length);
        Add(x2Local, x2Local, y2Local, length);

        // Duplicate(tLocal, xyz1Gm.GetValue(idx + 2), length);
        // DumpTensor(tLocal, 5, 16);
        // Muls(tLocal, tLocal, (float)0.0, length);
        // DumpTensor(tLocal, 5, 16);
        // Adds(tLocal, tLocal, (float)7.0, length);
    

        // for (int i = 0; i < length; i++)
        // {
        //     tLocal.SetValue(i, xyz1Gm.GetValue(idx + 2));
        // }
        Muls(tLocal, tLocal, (float)0.0, length);
        Adds(tLocal, tLocal, xyz1Gm.GetValue(idx + 2), length);
        Sub(z2Local, z2Local, tLocal, length);
        Mul(z2Local, z2Local, z2Local, length);
        Add(x2Local, x2Local, z2Local, length);
        
        // DataCopy(dGm[progress * 8], x2Local, 16);

        uint32_t val = 0;
        if (progress == tileNum - 1)
        {
            val = m - progress * tileDataNum / 3;
        }
        else
        {
            val = tileDataNum / 3;
        }

        for (int32_t i = 0; i < val; i++)
        {
            if (x2Local.GetValue(i) < bLocal.GetValue(0))
            {
                bLocal.SetValue(2, bLocal.GetValue(1));
                biLocal.SetValue(2, biLocal.GetValue(1));
                bLocal.SetValue(1, bLocal.GetValue(0));
                biLocal.SetValue(1, biLocal.GetValue(0));
                bLocal.SetValue(0, x2Local.GetValue(i));
                biLocal.SetValue(0, (int32_t)progress * tileDataNum / 3 + i);
            }
            else if (x2Local.GetValue(i) < bLocal.GetValue(1))
            {
                bLocal.SetValue(2, bLocal.GetValue(1));
                biLocal.SetValue(2, biLocal.GetValue(1));
                bLocal.SetValue(1, x2Local.GetValue(i));
                biLocal.SetValue(1, (int32_t)progress * tileDataNum / 3 + i);
            }
            else if (x2Local.GetValue(i) < bLocal.GetValue(2))
            {
                bLocal.SetValue(2, x2Local.GetValue(i));
                biLocal.SetValue(2, (int32_t)progress * tileDataNum / 3 + i);
            }
        }

        // dLocal.SetValue(idx + 0, bLocal.GetValue(0));
        // dLocal.SetValue(idx + 1, bLocal.GetValue(1));
        // dLocal.SetValue(idx + 2, bLocal.GetValue(2));
        // iLocal.SetValue(idx + 0, biLocal.GetValue(0));
        // iLocal.SetValue(idx + 1, biLocal.GetValue(1));
        // iLocal.SetValue(idx + 2, biLocal.GetValue(2));

        // DataCopy(dGm[idx], bLocal, 8);
        // DataCopy(iGm[idx], biLocal, 8);
        dGm.SetValue(idx + 0, bLocal.GetValue(0));
        dGm.SetValue(idx + 1, bLocal.GetValue(1));
        dGm.SetValue(idx + 2, bLocal.GetValue(2));
        iGm.SetValue(idx + 0, biLocal.GetValue(0));
        iGm.SetValue(idx + 1, biLocal.GetValue(1));
        iGm.SetValue(idx + 2, biLocal.GetValue(2));
        // DataCopy(dGm, bLocal, 16);
        // DataCopy(iGm, biLocal, 16);


        x2Queue.FreeTensor(x2Local);
        y2Queue.FreeTensor(y2Local);
        z2Queue.FreeTensor(z2Local);
        // bestQueue.FreeTensor(bLocal);
        // bestiQueue.FreeTensor(biLocal);

        // dQueue.EnQue<DTYPE_DIST>(dLocal);
        // iQueue.EnQue<DTYPE_INDICES>(iLocal);

    }

    __aicore__ inline void CopyOut(int32_t batchIdx, int32_t progress, uint32_t length)
    {
        LocalTensor<DTYPE_DIST> dLocal = dQueue.DeQue<DTYPE_DIST>();
        LocalTensor<DTYPE_INDICES> iLocal = iQueue.DeQue<DTYPE_INDICES>();

        // DataCopy(dGm[batchIdx * n * 3 + progress * tileDataNum], dLocal, length);
        // DataCopy(iGm[batchIdx * n * 3 + progress * tileDataNum], iLocal, length);

        dQueue.FreeTensor(dLocal);
        iQueue.FreeTensor(iLocal);

    }

    __aicore__ inline void InitBest()
    {
        LocalTensor<DTYPE_XYZ1> bLocal = bestQueue.AllocTensor<DTYPE_XYZ1>();
        LocalTensor<DTYPE_INDICES> biLocal = bestiQueue.AllocTensor<DTYPE_INDICES>();

        bLocal.SetValue(0, (DTYPE_XYZ1)1e38);
        bLocal.SetValue(1, (DTYPE_XYZ1)1e38);
        bLocal.SetValue(2, (DTYPE_XYZ1)1e38);
        biLocal.SetValue(0, (DTYPE_INDICES)0);
        biLocal.SetValue(1, (DTYPE_INDICES)0);
        biLocal.SetValue(2, (DTYPE_INDICES)0);

        bestQueue.EnQue(bLocal);
        bestiQueue.EnQue(biLocal);
    }

    __aicore__ inline void InitBestBuf()
    {
        LocalTensor<DTYPE_XYZ1> bLocal = bestBuf.Get<DTYPE_XYZ1>();
        LocalTensor<DTYPE_INDICES> biLocal = bestiBuf.Get<DTYPE_INDICES>();

        bLocal.SetValue(0, (DTYPE_XYZ1)FLT_MAX);
        bLocal.SetValue(1, (DTYPE_XYZ1)FLT_MAX);
        bLocal.SetValue(2, (DTYPE_XYZ1)FLT_MAX);
        // bLocal.SetValue(0, (DTYPE_XYZ1)1e38);
        // bLocal.SetValue(1, (DTYPE_XYZ1)1e38);
        // bLocal.SetValue(2, (DTYPE_XYZ1)1e38);
        biLocal.SetValue(0, (DTYPE_INDICES)0);
        biLocal.SetValue(1, (DTYPE_INDICES)0);
        biLocal.SetValue(2, (DTYPE_INDICES)0);
    }



    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> xyz0Queue, xyz1Queue;
    TQue<QuePosition::VECIN, BUFFER_NUM> x2Queue, y2Queue, z2Queue;
    TQue<QuePosition::VECIN, BUFFER_NUM> bestQueue, bestiQueue;
    TQue<QuePosition::VECOUT, BUFFER_NUM> dQueue, iQueue;
    TBuf<QuePosition::VECCALC> bestBuf, bestiBuf;
    TBuf<QuePosition::VECCALC> tmpBuffer;

    GlobalTensor<DTYPE_XYZ1> xyz1Gm, xyz2Gm;
    GlobalTensor<DTYPE_DIST> dGm;
    GlobalTensor<DTYPE_INDICES> iGm;


    uint32_t tileNum = 0;
    uint32_t coreDataNum = 0;
    uint32_t tileDataNum = 0;
    uint32_t startPointer = 0;
    uint32_t b = 0;
    uint32_t m = 0;
    uint32_t n = 0;
    uint32_t batchDataNum = 0;
};

extern "C" __global__ __aicore__ void three_nn(GM_ADDR xyz1, GM_ADDR xyz2, GM_ADDR dist, GM_ADDR indices, GM_ADDR workspace, GM_ADDR tiling) {
    
    GET_TILING_DATA(tiling_data, tiling);
    
    KernelThreeNN op;
    op.Init(xyz1, xyz2, dist, indices,
            tiling_data.ALIGN_NUM, tiling_data.ub_data_num,
            tiling_data.b, tiling_data.m, tiling_data.n, 
            tiling_data.batchDataNum, tiling_data.batchDataNumAligned,
            tiling_data.core_size, tiling_data.core_remain);
    op.Process();
    // op.test();
}